/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file conv2d_api_impl.h
 * \brief conv2d api impl class
 */

#ifndef CONV2D_API_IMPL_H
#define CONV2D_API_IMPL_H

#include "kernel_utils.h"
#include "conv2d_framework_util.h"
#include "conv2d_config.h"
#include "conv2d_common_func.h"
#include "conv2d_common_sub_api.h"

using namespace AscendC;

namespace conv2d {
template<typename Intf, class ConfigIn>
struct Conv2dApiImpl {
public:
    using Config = ConfigIn;
    using ConvParam = typename ConfigIn::ConvParam;
    constexpr static uint32_t ImplType = Config::implType;

public:
    __aicore__ inline Conv2dApiImpl() {}

    CONV_REG_IMPL(Config, Conv2dFunc, Init);
    CONV_REG_IMPL(Config, Conv2dFunc, SetFmap);
    CONV_REG_IMPL(Config, Conv2dFunc, SetWeight);
    CONV_REG_IMPL(Config, Conv2dFunc, SetScale);
    CONV_REG_IMPL(Config, Conv2dFunc, SetBias);
    CONV_REG_IMPL(Config, Conv2dFunc, SetOrgFmapShape);
    CONV_REG_IMPL(Config, Conv2dFunc, SetOrgWeightShape);
    CONV_REG_IMPL(Config, Conv2dFunc, SetOrgOutputShape);
    CONV_REG_IMPL(Config, Conv2dFunc, SetSingleFmapShape);
    CONV_REG_IMPL(Config, Conv2dFunc, SetSingleWeightShape);
    CONV_REG_IMPL(Config, Conv2dFunc, SetSingleOutputShape);
    CONV_REG_IMPL(Config, Conv2dFunc, SetFmapStartPosition);
    CONV_REG_IMPL(Config, Conv2dFunc, Iterate);
    CONV_REG_IMPL(Config, Conv2dFunc, IterateAll);
    CONV_REG_IMPL(Config, Conv2dFunc, GetTensorC);
    CONV_REG_IMPL(Config, Conv2dFunc, End);

    struct ContextData: public Config::ContextData {
        __aicore__ inline ContextData(){};
        TPipe pipe;
        const struct TConv2DTiling* __restrict conv2dTiling;
        // GM Tensor
        GlobalTensor<typename Config::FmapT> agm;
        GlobalTensor<typename Config::WeightT> bgm;
        GlobalTensor<typename Config::BiasT> biasgm;
        GlobalTensor<typename Config::ScaleT> scalegm;

        // LocalTensor
        LocalTensor<typename Config::FmapT> al1;
        LocalTensor<typename Config::WeightT> bl1;
        LocalTensor<typename Config::BiasT> biasL1;
        LocalTensor<typename Config::ScaleT> scaleL1;
        LocalTensor<typename Config::L0cT> biasBT;
        LocalTensor<typename Config::FmapT> al0;
        LocalTensor<typename Config::WeightT> bl0;
        LocalTensor<typename Config::L0cT> cl0;

        // Queue
        TQue<QuePosition::A1, 1> queueAL1; // AL1
        TQue<QuePosition::B1, 1> queueBL1; // BL1
        TQue<QuePosition::A1, 1> queueBiasL1; // BiasL1
        TQue<TPosition::C2, 1> queueBiasBT; // BT
        TQue<QuePosition::A1, 1> queueScaleL1; // ScaleL1
        TQue<QuePosition::CO2, 1> queueOutput; // Cgm
        TQue<QuePosition::A2, 1> queueAL0; // AL0
        TQue<QuePosition::B2, 1> queueBL0; // BL0
        TQue<QuePosition::CO1, 1> queueCL0; // CL0

        uint8_t enableBias = false;  // 是否有bias
        uint8_t enableScale = false;  // 是否有scale
        uint8_t isFirstIterate = true; // 是否第一次Iterate
        uint8_t loadAL1Flag = true; // 是否载入AL1的标志
        uint8_t loadBL1Flag = true; // 是否载入BL1的标志
        uint8_t loadAL0Flag = true; // 是否载入AL0的标志
        uint8_t loadBL0Flag = true; // 是否载入BL0的标志

        uint64_t oriBatch = 0; //  fmap上原始batch大小
        uint64_t oriCi = 0;   //  fmap上cin大小
        uint64_t oriCo = 0;  //  weight上cout大小
        uint64_t oriHi = 0;  //  fmap上h大小
        uint64_t oriWi = 0;  //  fmap上w大小
        uint64_t oriHo = 0;  //  output上h大小
        uint64_t oriWo = 0;  //  output上w大小
        uint64_t kernelH = 0; //  weight上h大小
        uint64_t kernelW = 0; //  weight上w大小
        uint64_t singleCoreCin = 0; // 单核上处理的Cin大小
        uint64_t singleCoreHi = 0; // 单核上处理的Hi大小
        uint64_t singleCoreWi = 0; // 单核上处理的Hi大小
        uint64_t singleCoreCo = 0; // 单核上处理的Co大小
        uint64_t singleCoreHo = 0; // 单核上处理的Ho大小
        uint64_t singleCoreWo = 0; // 单核上处理的Wo大小

        uint64_t kAL1Iter = 0;    // AL1上k方向迭代器
        uint64_t hoAL1Iter = 0;  // AL1上ho方向迭代器
        uint64_t woAL1Iter = 0;  // AL1上wo方向迭代器
        uint64_t kBL1Iter = 0;    // BL1上k方向迭代器
        uint64_t nBL1Iter = 0;   // BL1上n方向迭代器
        uint64_t kIter = 0; // k方向迭代器，从DDR到L0
        uint64_t kL0AIter = 0; // L1A 到L0方向的迭代器
        uint64_t kL0BIter = 0; // L1B 到L0方向的迭代器
        uint64_t nL0Iter = 0; // BL0上n方向迭代器
        uint64_t woL0Iter = 0; // L0A上wo方向迭代器
        uint64_t hoL0Iter = 0; // L0A上ho方向迭代器

        uint64_t kAL1Tail = 0;  // AL1上k方向的尾块大小
        uint64_t hoAL1Tail = 0; // AL1上ho方向的尾块大小
        uint64_t woAL1Tail = 0; // AL1上wo方向的尾块大小
        uint64_t kBL1Tail = 0; // BL1上k方向的尾块大小
        uint64_t nBL1Tail = 0; // BL1上N方向的尾块大小
        uint64_t nL0Tail = 0; // BL0上n方向的对齐后尾块大小
        uint64_t nL0TailNotAlign = 0; // BL0上n方向的真实尾块大小
        uint64_t kL0Tail = 0;
        uint64_t woL0Tail = 0;
        uint64_t hoL0Tail = 0;
        uint64_t woL1SmallTail = 0; // wo方向16不对齐场景会存在小尾块

        Conv2dFunc::MMadTools<Intf>  madIns;
        Conv2dFunc::LoadBL0Tools<Intf> loadBL0Ins;
        Conv2dFunc::LoadBL1Tools<Intf> loadBL1Ins;
        Conv2dFunc::LoadAL1Tools<Intf> loadAl1Ins;
        Conv2dFunc::LoadAL0Tools<Intf> loadAL0Ins;
        Conv2dFunc::LoadChannelWiseL1Tools<Intf, typename Config::BiasT> loadBiasL1Ins;
        Conv2dFunc::LoadBiasBtTools<Intf> loadBiasBTIns;
        Conv2dFunc::LoadChannelWiseL1Tools<Intf, typename Config::ScaleT> loadScaleL1Ins;
        Conv2dFunc::CopyOutTools<Intf> copyOutIns;

        // Iterate时需要用到的变量
        size_t maxKAL1Iter = 0;
        size_t maxHoL1Iter = 0;
        size_t maxWoL1Iter = 0;
        size_t maxNBL1Iter = 0;
        size_t maxKBL1Iter = 0;
        size_t maxNL0Iter = 0;
        size_t maxKL0Iter = 0;
        size_t maxWoL0Iter = 0;
        size_t maxHoL0Iter = 0;
  
        uint64_t ddr2l1LoopN = 0;
        uint64_t l12l0LoopN = 0;
        uint64_t ddr2l1LoopW = 0;
        uint64_t ddr2l1LoopH = 0;
        uint64_t ddr2l0LoopK = 0;
        uint64_t l12l0LoopW = 0;
        uint64_t l12l0LoopH = 0;

        uint64_t currentWoL1 = 0;
        uint64_t currentHoL1 = 0;
        uint64_t currentWoL0 = 0;
        uint64_t currentHoL0 = 0;
        uint64_t kernelHxkernelW = 0;
        uint64_t dilatedKernelHxkernelW = 0;
        uint64_t dilatedKernelH = 0;
        uint64_t dilatedKernelW = 0;

        uint64_t multiKAL1 = 1;
        uint64_t multiKBL1 = 1;
        bool kAL1fullload = false;
        bool kBL1fullload = false;
        bool biasFullLoadFlag = false;
        bool scaleFullLoadFlag = false;
        int64_t singleCoreHiStartPos = -1;
        int64_t wiStartPos = 0;
        uint64_t hoL1 = 1;
        uint64_t woL1 = 16;
    };

    // 各个模板函数的变量，暂未使用
    struct ImplDataType {
        __aicore__ inline ImplDataType() {};
    };
};
}  // namespace conv2d

#endif // __CONV2D_API_IMPL_H__

